#pragma once
#include <iostream>
#include <random>
#include <vector>

class CSCMatrix {
  public:
    CSCMatrix(int rowCount, int colCount) {
        construct(rowCount, colCount);
    }

    std::vector<long long> colOffset_;
    std::vector<long long> rowInd_;
    std::vector<double> data_;
    int rowCount_, colCount_;

    double get(int row, int col) const {
        int currRow = -1;

        for (int pos = colOffset_[col]; pos < colOffset_[col + 1]; ++pos) {
            currRow = rowInd_[pos];

            if (currRow == row) {
                return data_[pos];

            } else if (currRow > row) {
                break;
            }
        }

        return 0;
    }
    CSCMatrix &set(double val, int row, int col) {

        int pos = colOffset_[col];
        int currRow = -1;

        for (; pos < colOffset_[col + 1]; pos++) {
            currRow = rowInd_[pos];

            if (currRow >= row) {
                break;
            }
        }

        if (currRow != row) {
            if (!(val == 0)) {
                insert(pos, row, col, val);
            }

        } else if (val == 0) {
            remove(pos, row);

        } else {
            data_[pos] = val;
        }

        return *this;
    }

  private:
    void insert(int index, int row, int col, double val) {

        data_.insert(data_.begin() + index, val);
        rowInd_.insert(rowInd_.begin() + index, row);

        for (int i = col + 1; i <= colCount_; i++) {
            colOffset_[i] += 1;
        }
    }

    void remove(int index, int col) {
        if (data_.size() > 0) {
            data_.erase(data_.begin() + index);
            rowInd_.erase(rowInd_.begin() + index);

            for (int i = col + 1; i <= colCount_; i++) {
                colOffset_[i] -= 1;
            }
        }
    }

    void construct(int rows, int columns) {
        if (rows < 1 || columns < 1) {
            std::cout << "Matrix dimensions cannot be zero or negative.\n";
            exit(1);
        }

        rowCount_ = rows;
        colCount_ = columns;

        // 矩阵下标范围是 [0,m-1] [0,n-1] 左闭右闭，三个数组的下标是从0开始的
        data_ = std::vector<double>();
        data_.reserve(rows * columns / 10);
        rowInd_ = std::vector<long long>();
        rowInd_.reserve(rows * columns / 10);
        colOffset_ = std::vector<long long>(columns + 1, 0);
    }

    friend std::ostream &operator<<(std::ostream &os, const CSCMatrix &matrix) {
        for (int i = 0; i < matrix.rowCount_; i++) {
            for (int j = 0; j < matrix.colCount_; j++) {
                if (j != 0) {
                    os << " ";
                }

                os << matrix.get(i, j);
            }

            if (i < matrix.rowCount_) {
                os << std::endl;
            }
        }
        return os;
    }
}; // class CSRMatrix

void csc_random_fill(CSCMatrix &matrix) {
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_int_distribution<> distrib(1, 5);
    std::uniform_int_distribution<> prob(0, 1000);
    for (int i = 0; i < matrix.rowCount_; i++) {
        for (int j = 0; j < matrix.colCount_; j++) {
            if (prob(gen) > 500) {
                double rand_num = distrib(gen);
                matrix.set(rand_num, i, j);
            }
        }
    }
}

void csc_zero_fill(CSCMatrix &matrix) {
    for (int i = 0; i < matrix.rowCount_; i++) {
        for (int j = 0; j < matrix.colCount_; j++) {
            matrix.set(0, i, j);
        }
    }
}
void csc_one_fill(CSCMatrix &matrix) {
    for (int i = 0; i < matrix.rowCount_; i++) {
        for (int j = 0; j < matrix.colCount_; j++) {
            matrix.set(1, i, j);
        }
    }
}
